# Copyright (c) 2023; NVIDIA CORPORATION. All rights reserved.
# Copyright (c) 2025, Huawei Technologies Co., Ltd.  All rights reserved.

from typing import Optional

import torch
from torch import Tensor

from megatron.core.transformer.transformer_config import TransformerConfig
from megatron.core.transformer.transformer_block import TransformerBlock
from megatron.core.inference.contexts import BaseInferenceContext
from megatron.core.packed_seq_params import PackedSeqParams

from megatron.core.utils import deprecate_inference_params


_ROTATION_MATRIX = None


def get_rotation_matrix(x):
    global _ROTATION_MATRIX
    if _ROTATION_MATRIX is None:
        import numpy as np
        dim = x.shape[-1]
        index1 = np.ones(dim)
        index1[::2] = 0
        index2 = np.zeros(dim)
        index2[::2] = -1
        rotation_matrix = np.eye(dim, k=1) * index1 + np.eye(dim, k=-1) * index2
        _ROTATION_MATRIX = (
            torch.from_numpy(rotation_matrix[None, None, :, :]).to(x.dtype).to(x.device)
        )
    return _ROTATION_MATRIX


def local_rotate_half(x: Tensor, rotary_interleaved: bool) -> Tensor:
    """Change sign so the last dimension becomes [-odd, +even]

    Args:
        x (Tensor): Input tensor

    Returns:
        Tensor: Tensor rotated half
    """
    if not rotary_interleaved:
        x1, x2 = torch.chunk(x, 2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)
    return torch.matmul(x, get_rotation_matrix(x))